from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt

def yolov8_inference(image_path, model_type='best.pt'):
    """
    使用YOLOv8模型进行目标检测
    
    参数:
        image_path: 输入图像的路径
        model_type: 预训练模型类型
    """
    # 加载预训练模型
    model = YOLO(model_type)
    
    # 进行推理
    results = model(image_path)
    
    # 处理结果
    for result in results:
        # 绘制检测结果
        annotated_img = result.plot()
        
        # 转换颜色通道（OpenCV默认BGR，Matplotlib需要RGB）
        annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
        
        # 显示结果
        plt.figure(figsize=(10, 10))
        plt.imshow(annotated_img_rgb)
        plt.axis('off')
        plt.show()
        
        # 打印检测到的目标信息
        print("检测到的目标:")
        for box in result.boxes:
            class_id = result.names[int(box.cls)]
            coordinates = box.xyxy[0].tolist()
            confidence = box.conf[0].item()
            print(f"- {class_id}: 置信度 {confidence:.2f}, 坐标 {coordinates}")
    
    return results

if __name__ == "__main__":
    # 示例：处理单张图像
    image_path = "/home/fourzkw/Project/yolo_test/yolo_test_py/1.jpg"  # 替换为你的图像路径
    model_path = "/home/fourzkw/Project/yolo_test/yolo_test_py/best.pt"
    yolov8_inference(image_path, model_path)